package top.kangyaocoding.tech.common;

import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * 描述: RedisChatMemory
 *
 * @author K·Herbert
 * @since 2025-04-13 22:56
 */
@Component
@RequiredArgsConstructor
public class RedisChatMemory implements ChatMemory {

    private static final String REDIS_KEY_PREFIX = "REDIS_CHAT_MEMORY:";
    // 默认过期时间为72小时
    private static final long DEFAULT_EXPIRATION_TIME_MS = 72 * 60 * 60 * 1000;
    private final RedisTemplate<String, Message> redisTemplate;

    @Override
    public void add(String conversationId, List<Message> messages) {
        if (conversationId == null || conversationId.isEmpty()) {
            throw new IllegalArgumentException("conversationId 不能为 null 或为空");
        }
        // 如果消息为空，直接返回
        if (messages == null || messages.isEmpty()) {
            return;
        }

        String key = REDIS_KEY_PREFIX + conversationId;
        try {
            redisTemplate.opsForList().rightPushAll(key, messages);
            redisTemplate.expire(key, DEFAULT_EXPIRATION_TIME_MS, TimeUnit.MILLISECONDS);
        } catch (Exception e) {
            throw new RuntimeException("无法向 Redis 添加消息", e);
        }
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        if (conversationId == null || conversationId.isEmpty()) {
            throw new IllegalArgumentException("conversationId 不能为 null 或为空");
        }
        if (lastN <= 0) {
            throw new IllegalArgumentException("lastN 必须大于 0");
        }

        String key = REDIS_KEY_PREFIX + conversationId;
        try {
            List<Message> serializedMessages = redisTemplate.opsForList().range(key, -lastN, -1);
            return serializedMessages != null ? serializedMessages : List.of();
        } catch (Exception e) {
            throw new RuntimeException("无法从 Redis 检索消息", e);
        }
    }

    @Override
    public void clear(String conversationId) {
        if (conversationId == null || conversationId.isEmpty()) {
            throw new IllegalArgumentException("conversationId 不能为 null 或为空");
        }

        String key = REDIS_KEY_PREFIX + conversationId;
        try {
            redisTemplate.delete(key);
        } catch (Exception e) {
            throw new RuntimeException("无法从 Redis 清除消息", e);
        }
    }
}

